#pragma once

#include <memory>
#include <vector>
#include "autograd/autograd.h"
#include "autograd/add.h"
#include "autograd/mul.h"
#include "autograd/sub.h"
#include "autograd/div.h"
#include "autograd/square.h"
#include "autograd/neg.h"

namespace  ldl
{
namespace autograd
{
template<typename T>
class Variable;

template<typename T>
class Function : public std::enable_shared_from_this<Function<T>>
{
public:
    std::vector<Variable<T>> operator()(std::vector<Variable<T>>& inputs) {
        m_inputs.clear();
        int64_t max_generation = 0;
        for(auto& input_variable:inputs) {
            m_inputs.push_back(input_variable.m_data_ptr());
            max_generation = std::max(input_variable.m_data_ptr()->generation(), max_generation);
        }
        m_generation = max_generation;

        if(m_first_operate) {
            m_outputs_weak_ptr.resize(m_outputs.size());

            // 转换 shared_ptr → weak_ptr
            m_outputs_weak_ptr.assign(m_outputs.begin(), m_outputs.end());
        }
        forward();

        std::vector<Variable<T>> ret{};
        for(auto& item:m_outputs_weak_ptr) {
            item.lock()->set_creator(this->shared_from_this());
            item.lock()->set_generation(this->m_generation + 1);
            // LogInfo() << "this->m_generation: " << this->m_generation;
            ret.push_back(item.lock());
        }

        if(m_first_operate) {
            // 批量置空 shared_ptr
            std::fill(m_outputs.begin(), m_outputs.end(), nullptr);
        }
        return ret;
    }

    Variable<T> operator()(Variable<T> input) {
        m_inputs[0] = input.data();

        forward();

        m_outputs[0].set_creator(this->shared_from_this());
    
        return m_outputs[0];
    }

    virtual void forward() = 0;
    virtual void do_backward() = 0;
    virtual void backward() {
        for(auto& item:m_inputs) {
            if(item->grad() == nullptr) {
                item->grad() = std::make_shared<Variable<T>>(zeros_like(*(item->data())));
            }
        }
        do_backward();
    }

    const std::vector<std::shared_ptr<Autograd<T>>>& inputs() const {
        return m_inputs;
    }

    const std::vector<std::weak_ptr<Autograd<T>>>& outputs() const {
        return m_outputs_weak_ptr;
    }

    static Variable<T> add(Variable<T> a, Variable<T> b) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Add<T>>();
        return operate(a, b, creator);
    }

    static Variable<T> sub(Variable<T> a, Variable<T> b) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Sub<T>>();
        return operate(a, b, creator);
    }

    static Variable<T> mul(Variable<T> a, Variable<T> b) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Mul<T>>();
        return operate(a, b, creator);
    }

    static Variable<T> div(Variable<T> a, Variable<T> b) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Div<T>>();
        return operate(a, b, creator);
    }

    static Variable<T> square(Variable<T> a) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Square<T>>();
        std::vector<Variable<T>> inputs{};
        inputs.push_back(a);
        return (*creator)(inputs)[0];
    }

    static Variable<T> neg(Variable<T> a) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Neg<T>>();
        std::vector<Variable<T>> inputs{};
        inputs.push_back(a);
        return (*creator)(inputs)[0];
    }

    static Variable<T> operate(Variable<T>& a, Variable<T>& b,  std::shared_ptr<Function<T>> creator) {
        std::vector<Variable<T>> inputs{};
        inputs.push_back(a);
        inputs.push_back(b);
        return (*creator)(inputs)[0];
    }

    int64_t generation() {
        return m_generation;
    }

    void set_generation(int64_t generation) {
        m_generation = generation;
    }

    virtual ~Function() = default;

    std::vector<std::shared_ptr<Autograd<T>>>& inputs() {
        return m_inputs;
    }
protected:
    std::vector<std::shared_ptr<Autograd<T>>> m_inputs;
    std::vector<std::weak_ptr<Autograd<T>>> m_outputs_weak_ptr;
    std::vector<std::shared_ptr<Autograd<T>>> m_outputs;
    int64_t m_generation{};
    bool m_first_operate{true};
};

// 比较Function的generation
struct CompareFunctionGeneration {
    template<typename T>
    bool operator()(const Function<T>& a, const Function<T>& b) const {
        return a.generation() < b.generation();
    }
};
}
}
